#include <webby/server.hxx>
#include <boost/bind.hpp>
#include <errno.h>
#include <assert.h>
#include <poll.h>
#include <unistd.h>
#include <fcntl.h>


#include <iostream>
#include <dlfcn.h>


#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include <webby/fastcgi.hxx>
#include <webby/logging.hxx>

using namespace webby;

Server::Server() : mCanExit( false ), mThreadsRunning( false ) , mServerSocket(0), mID(1)
{
	FD_ZERO( &read_set );
	FD_ZERO( &write_set );
	FD_ZERO( &except_set );

	addFastCGIHandler( "php", 2000 );
}

void 
Server::handleSocketClose( int clientSocket )
{
	boost::mutex::scoped_lock lock( mLocalLock );
	mSocketCloseSet.insert( clientSocket );
	fcntl( clientSocket, F_SETFL, O_NONBLOCK );

	WEBBY_LOG( Logging::DEBUG, "Adding socket handle for future close" );
}

void
Server::start( unsigned int port, unsigned int threads ) 
{
	WEBBY_LOG( Logging::INFORMATIONAL, "Starting Webby Server" );
	mServerSocket = socket( AF_INET, SOCK_STREAM,  IPPROTO_TCP );
	if ( mServerSocket  != -1 ) 
	{
		int i = 1;
		setsockopt( mServerSocket, SOL_SOCKET, SO_REUSEADDR, &i, sizeof(i) );
		sockaddr_in addr;
		memset( &addr, 0, sizeof(addr) );
		addr.sin_port = htons( port );		
		addr.sin_family = AF_INET;
		addr.sin_addr.s_addr = htonl( INADDR_ANY );

		if ( bind( mServerSocket, (sockaddr *)&addr, sizeof(addr) ) == 0 )
		{
			if ( listen( mServerSocket, 50 ) != -1 )
			{
				WEBBY_LOG( Logging::INFORMATIONAL, "Now listening on port " << port );

				addModule( "basemodule.so" );
				addModule( "cachemodule.so" );
			
				addFdToPollMap( mServerSocket );
			}
		} else {
			WEBBY_LOG( Logging::CRITICAL, "Unable to bind to port " << port );
			exit( -1 );
		}
	} else {
		WEBBY_LOG( Logging::CRITICAL, "Unable to create the main server socket" );
		exit( -1 );
	}

	// Starting threads
	WEBBY_LOG( Logging::DEBUG, "Starting " << threads << " worker threads" );
	for ( unsigned int i = 0 ; i < threads; ++i)
	{
		thread_ptr thread = thread_ptr( new boost::thread( boost::bind( &Server::threadStarted, this, i+1 ) ) );
		mThreads.push_back( thread );
	}
}

bool
Server::isRunning() const
{
	return !mCanExit;
}

void
Server::threadStarted( unsigned int threadNum ) 
{
	WEBBY_LOG( Logging::DEBUG, "Starting thread " << threadNum );

	while ( !mCanExit )
	{
		message_ptr message;	
		{
			boost::mutex::scoped_lock lock(mLock);
			mCondition.wait( lock );

			// we have the lock here
			if ( mThreadQueue.empty() ) continue;

			message = mThreadQueue.front();
			mThreadQueue.pop();
		}

		assert( message.get() );
		WEBBY_LOG( Logging::DEBUG, "About to perform primary task on thread " << threadNum );
		message->doWork();
		WEBBY_LOG( Logging::DEBUG, "Task complete on thread " << threadNum );

		bool empty = false;
		while ( !empty )
		{
			// scope this lock so we don't hold it during exuction of the message
			{
				boost::mutex::scoped_lock lock(mLock);
				if ( !mThreadQueue.empty() )
				{
					message = mThreadQueue.front();
					mThreadQueue.pop();
				}
				else
				{
					message.reset();
					empty = true;
				}
			}

			if ( message.get() != 0) 
			{
				WEBBY_LOG( Logging::DEBUG, "About to perform secondary task on thread " << threadNum );
				message->doWork();
				WEBBY_LOG( Logging::DEBUG, "Secondary tast complete on thread " << threadNum );
			}
		}
	}

	assert( mCanExit );
	WEBBY_LOG( Logging::DEBUG, "Exiting thread " << threadNum );
}

void
Server::handleEvents() 
{
	std::map< int, pollfd > pollMap;

	FD_ZERO( &read_set );
	FD_ZERO( &write_set );
	FD_ZERO( &except_set );

	FD_SET( mServerSocket, &read_set );
	FD_SET( mServerSocket, &except_set );

	for ( std::set< unsigned int >::const_iterator iter = mSockets.begin(); iter != mSockets.end(); ++iter)
	{
		FD_SET( *iter, &read_set );
		FD_SET( *iter, &except_set );
	}

	for ( std::map< int, FastCGITransaction >::const_iterator iter = mSocketToFastCGIMap.begin(); iter != mSocketToFastCGIMap.end(); ++iter)
	{
		//std::cout << "Adding FastCGI Socket " << iter->first << std::endl;
		FD_SET( iter->first, &read_set );
		FD_SET( iter->first, &except_set );
	}

	unsigned int maxSocket = mServerSocket;

	timeval tv;
	tv.tv_sec = 0;
	tv.tv_usec = WEBBY_CONFIG_SELECT_TIME;

	if ( !mSockets.empty() ) {
		unsigned int maxClientSocket = *mSockets.rbegin();
		maxSocket = std::max< unsigned int >( maxSocket, maxClientSocket );
	}

	if ( !mSocketToFastCGIMap.empty() ) 
	{
		unsigned int maxFastCGISocket = (mSocketToFastCGIMap.rbegin())->first;
		maxSocket = std::max< unsigned int >( maxSocket, maxFastCGISocket );
	}


	int result = select( maxSocket + 1, &read_set, &write_set, &except_set, &tv );
	if ( result == -1 ) 
	{
		// error
	} else if ( result ) {
		// new data
		//std::cout << "New data..." << std::endl;
		if ( FD_ISSET( mServerSocket, &read_set ) ) 
		{
			handleServerAccept();		
		}

		if ( FD_ISSET( mServerSocket, &except_set ) ) 
		{
			std::cout << "Exception on main socket" << std::endl;
		}


		std::set< unsigned int > removeSockets;
		for ( std::map< int, FastCGITransaction >::const_iterator iter = mSocketToFastCGIMap.begin(); iter != mSocketToFastCGIMap.end(); ++iter)
		{
			if ( FD_ISSET( iter->first, &read_set ) ) 
			{
				//std::cout << "Found FastCGI Input" << std::endl;
				bool erase = false;
				handleFastCGISocket( iter->first, erase );

				if ( erase )
				{
					removeSockets.insert( iter->first );
				}
			}
		}

		for ( std::set< unsigned int >::const_iterator iter = removeSockets.begin(); iter != removeSockets.end(); ++iter)
		{
			mSocketToFastCGIMap.erase( *iter );
		}

		removeSockets.clear();
		{
			for ( std::set< unsigned int >::const_iterator iter = mSockets.begin(); iter != mSockets.end(); ++iter)
			{
				handleClientRead( *iter );
				removeSockets.insert( *iter );
			}
		}

		{
			// remove sockets from the list -- inefficient
			for ( std::set< unsigned int >::const_iterator iter = removeSockets.begin(); iter != removeSockets.end(); ++iter)
			{
				mSockets.erase( *iter );
			}
		}

	
		

	}

	{
		boost::mutex::scoped_lock lock( mLocalLock );
		std::set< int > closeSockets = mSocketCloseSet;
		for ( std::set< int>::const_iterator iter = closeSockets.begin(); iter != closeSockets.end(); ++iter) 
		{
			char buffer[1024];
			size_t amountRead = recv( (*iter), buffer, sizeof( buffer ), 0 );
			if ( amountRead == 0 ) 
			{
				WEBBY_LOG( Logging::DEBUG, "Previous request socket closed" );
				mSocketCloseSet.erase( *iter );

				close( *iter );
			}
		}
	}


	while (true)
	{
		boost::shared_ptr< Message > messagePtr;
		{
			{
				boost::mutex::scoped_lock lock( mLocalLock );
				if ( !mLocalThreadQueue.empty() )
				{
					messagePtr = mLocalThreadQueue.front();
					mLocalThreadQueue.pop();
				}
			}

			if ( messagePtr.get() == 0)
			{
				break;
			}
	
			messagePtr->doWork();
		}
	}
}

void
Server::handleServerAccept() 
{
	sockaddr_in addr;
	socklen_t len = sizeof( addr );
	//std::cout << "Accepting on " << mServerSocket << std::endl;
	int s = accept( mServerSocket, (sockaddr *)&addr, &len );
	if ( s != -1 ) 
	{
#if defined(WEBBY_DEBUG)
		std::cout << "New connection on " << s << std::endl;
#endif
		// let's go non blocking
		//fcntl( s, F_SETFL, O_NONBLOCK );     
		int i = 64*1024;
		setsockopt (s, SOL_SOCKET, SO_SNDBUF, &i, sizeof(i) );
		setsockopt (s, SOL_SOCKET, SO_RCVBUF, &i, sizeof(i) );

		mSockets.insert( s );
	} else {
		std::cout << "Main socket error " << errno << std::endl;
		exit( -1 );
	
	}
}

void 
Server::addToThreadQueue( message_ptr ptr )
{
	boost::mutex::scoped_lock lock(mLock);
	mThreadQueue.push( ptr );
	mCondition.notify_one();
}


void 
Server::addToLocalThreadQueue( message_ptr ptr )
{
	boost::mutex::scoped_lock lock(mLocalLock);
	mLocalThreadQueue.push( ptr );
}


void
Server::handleClientRead( int readSocket )
{
	addToThreadQueue( message_ptr( new ReadSocketMessage( shared_from_this(), readSocket ) ) );
}

void 
Server::quit()
{
	WEBBY_LOG( Logging::DEBUG, "Quit requested" );
	mCanExit = true;
	mCondition.notify_all();

	WEBBY_LOG( Logging::DEBUG, "Terminating all worker threads" );
	for ( std::vector< thread_ptr >::iterator iter = mThreads.begin(); iter != mThreads.end(); ++iter)
	{
		(*iter)->join();
	}

	WEBBY_LOG( Logging::DEBUG, "Shutting down and closing all sockets" );
	for ( std::set< unsigned int >::const_iterator iter = mSockets.begin(); iter != mSockets.end(); ++iter)
	{
		shutdown( *iter, 2 );
		close( *iter );
	}

	shutdown( mServerSocket, 2 );
	close( mServerSocket );	

	WEBBY_LOG( Logging::INFORMATIONAL, "Web server successfully shutdown" );
}

void 
Server::addFastCGIHandler( const std::string &extension, unsigned short port )
{
	mFastCGIHandlers[extension] = FastCGIServiceInfo( port );
}

bool 
Server::isFastCGI( const std::string &extension ) const
{
	return ( mFastCGIHandlers.find( extension ) != mFastCGIHandlers.end() );
}

void 
Server::handleFastCGIScript( const std::string &cgiType, const std::string &scriptName, int clientSocket )
{
	std::map< std::string, FastCGIServiceInfo >::iterator iter;
	std::map< std::string, FastCGIServiceInfo >::iterator end_iter;
	{
		iter = mFastCGIHandlers.find( cgiType );
		end_iter = mFastCGIHandlers.end();
	}

	if ( iter != end_iter )
	{
		unsigned short port = iter->second.mPort;
		sockaddr_in addr;
		addr.sin_family = AF_INET;
		addr.sin_port = htons( port );
		addr.sin_addr.s_addr = inet_addr("127.0.0.1");

		int s = socket( AF_INET, SOCK_STREAM, IPPROTO_TCP );
		if ( s != -1 )
		{
			if ( connect( s, (sockaddr *)&addr, sizeof( addr ) ) != -1 )
			{
				//std::cout << "Connected To Fast CGI" << std::endl;

				FastCGITransaction transaction;			
				transaction.mClientSocket = clientSocket;
	
				FastCGI f( mID );
				mID++;

				f.makeBeginRequest();
				send( s, f.mBuffer, 16, 0 );
	
				f.encodeParam( "SCRIPT_FILENAME", scriptName );
				send( s, f.mBuffer, f.mBuffer[5] + 8, 0 );

				f.encodeParam( "", "" );
				send( s, f.mBuffer, f.mBuffer[5] + 8, 0 );

				f.makeStdIn();
				send( s, f.mBuffer, 8, 0 );

				//std::cout << "Made contact with server" << std::endl;

				{
					mSocketToFastCGIMap[s] = transaction;
				}

				return;
			} 
			else
			{
				std::cout << "Unable to connect to FastCGI server for " << cgiType << std::endl;
			}
		}
	} 

	assert(0);	// we don't have a handler, something is borked
	close( clientSocket );
}


void
Server::handleFastCGISocket( int socket, bool &erase )
{
	erase = false;
	char data[32*1024];
	//std::cout << "About To Read" << std::endl;
	unsigned int amountRead = recv( socket, data, 32*1024 - 1, 0 );	
	//std::cout << "Read " << amountRead << " bytes " << std::endl;

	std::map< int, FastCGITransaction>::iterator iter = mSocketToFastCGIMap.find( socket );
	if ( iter != mSocketToFastCGIMap.end() )
	{
		FastCGITransaction &transaction = iter->second;

		memcpy( &transaction.mBuffer[transaction.mBufferSize], data, amountRead );
		transaction.mBufferSize += amountRead;

		if ( amountRead == 0)
		{
			//std::cout << "Closing socket" << std::endl;
			// shutdown PHP socket
			shutdown( socket, 2 );
			close( socket );

			// socket has closed
			int clientSocket = transaction.mClientSocket;
			//std::string content = transaction.mContent;
			//std::string content;

			erase = true;

			char *newBuffer = new char[transaction.mBufferSize];
			memcpy( newBuffer, transaction.mBuffer, transaction.mBufferSize );

			addToThreadQueue( message_ptr( new FastCGIOutputMessage( shared_from_this(), newBuffer, transaction.mBufferSize, clientSocket ) ) );	
		}
	} else {
		assert( 0 );
		close( socket );
	}
}

void
Server::addModule( const std::string &moduleName )
{
	std::string module = "modules/" + moduleName;
	void *shared = dlopen( module.c_str(), RTLD_LAZY);
	if ( shared )
	{
		void *(*entry)() = (void *(*)())dlsym( shared, "webby_create_module" );
		if ( entry )
		{
			ModulePtr modulePtr = ModulePtr( static_cast< Module *>( entry() ) );
			mModules.push_back( modulePtr );

			WEBBY_LOG( Logging::INFORMATIONAL, "Module \"" << modulePtr->getModuleName() << "\" loaded" );		
		}
		else
		{
			WEBBY_LOG( Logging::WARNING, "Unable to locate entry point in module" );
		}
	}
	else
	{
		WEBBY_LOG( Logging::WARNING, "Unable to load module " << module );
	}
}

